import numpy as np
from cvfo_mask import CvFo
from tqdm import tqdm
import paddle
import pandas as pd

block_size = 8
voc = pd.read_pickle("voc.pandas_pickle")
data = pd.read_pickle("dataset.pandas_pickle")
model = CvFo(len(voc), 256, 6, block_size)
t = paddle.to_tensor(
    [voc.loc[voc["voc"] == "<|p_{}|>".format(i)].values[0][1:-1].astype("int") for i in range(0, 128)]).T.unsqueeze(0)
model.load_dict(paddle.load("cvfo_model.pdparams"))
# model.load_dict(paddle.load("cvfo_model.pdparams"))
opt = paddle.optimizer.Adam(learning_rate=0.00005, parameters=model.parameters())
loss_func = paddle.nn.CrossEntropyLoss()
loss_func1 = paddle.nn.MSELoss()
epochs = 200
batch_size = 300
epoch = list(range(0, len(data), batch_size))
bar = tqdm(epoch * epochs)
for i in bar:
    if i == 0:
        np.random.shuffle(data)
    batch_data = paddle.to_tensor(data[i:i + batch_size]).astype('int64')
    labels = batch_data[:, -1:, 1:]
    out,out1,out2 = model(batch_data[:, :-1, :-1],t)

    loss = loss_func(out, labels.squeeze(1))
    loss+= loss_func1(out1, out2)
    bar.set_description("loss:%.4f" % loss.item())
    opt.clear_grad()
    loss.backward()
    opt.step()

    if i == epoch[-1]:
        paddle.save(model.state_dict(), "cvfo_model.pdparams")
paddle.save(model.state_dict(), "cvfo_model.pdparams")
